# -*- coding: utf-8 -*-
# @time: 2024/6/6 15:10
# @file: conversion_premium_modified
# @author: tyshixi08
import pandas as pd

from get_data.origin_data import *

def premium_modified_cal(code = code_ls(), start_date = month_ls()[0].replace('-','') + '01', end_date = month_ls()[-1].replace('-','') + '01'):

    # 获取转股溢价率和转股价值
    df = convertible.get_indicators(code, start_date=start_date, end_date=end_date,fields=['conversion_premium', 'conversion_value']).reset_index()
    df = df.dropna()
    df = df[df['conversion_value'] != 0]

    df['conversion_value_reci'] = 1 / df['conversion_value']

    # 按日期分类
    grouped = df.groupby('date')

    # 对每一类数据进行线性拟合
    fits = {}
    for name, group in grouped:
        x = group['conversion_value_reci'].values
        y = group['conversion_premium'].values
        fit = np.polyfit(x, y, 1)  # 进行一次多项式拟合，即线性拟合
        fits[name] = fit

    df_params = pd.DataFrame.from_dict(fits, orient='index').reset_index().rename(columns={'index':'date', 0:'k', 1:'b'})

    df = pd.merge(df, df_params, how='outer', on='date')

    df['conversion_premium_expected'] = df['conversion_value_reci'] * df['k'] + df['b']

    df['premium_modified'] = df['conversion_premium'] - df['conversion_premium_expected']

    return df

# 获得修正转股溢价率中位数
def premium_modified_median():
    df = premium_modified_cal()
    df_covpremium_modified = df.groupby('date')['premium_modified'].median().reset_index().rename(columns={'premium_modified':'covpremium_modified_mid'})
    return df_covpremium_modified.sort_values('date')

# 数据存档
def covpremium_modified_mid_save():
    excel_file_path = 'covpremium_modified_mid.csv'
    #if os.path.exists('原始数据/' + excel_file_path):
    #    df = pd.read_csv('原始数据/' + excel_file_path)
    #    return df
    #else:
    save_CSV(premium_modified_median(), 'get_data/原始数据/' + excel_file_path.split('.')[0])
    df = pd.read_csv('get_data/原始数据/' + excel_file_path)
    return df

